#include "ompl/base/spaces/SE3StateSpace.h"
#include "ompl/tools/benchmark/Benchmark.h"
#include "ompl/geometric/SimpleSetup.h"
#include "ompl/geometric/planners/rrt/RRT.h"
#include "ompl/contrib/hplanner/hPRM.h"
#include "ompl/geometric/planners/kpiece/KPIECE1.h"
#include "omplapp/apps/AppBase.h"
#include "omplapp/config.h"
#include "omplapp/apps/SE3RigidBodyPlanning.h"

#include <vector>
using std::vector;

namespace ob = ompl::base;
namespace og = ompl::geometric;
namespace app = ompl::app;

void test (int tries, vector <vector<double> > &times, vector <int> &attempts);

void configureApartmentProblem (app::SE3RigidBodyPlanning &setup);
void configureEasyProblem (app::SE3RigidBodyPlanning &setup);
void configureCubiclesProblem (app::SE3RigidBodyPlanning &setup);

#define NR_TRIES 10

int main (int argc, char **argv)
{
    // User can supply number of tries as 2nd command line argument.  Otherwise, use default NR_TRIES.
    int nr_tries;
    if (argc == 2)
    {
        nr_tries = atoi (argv[1]);

        // make sure command line input is valid
        if (nr_tries == INT_MAX || nr_tries == INT_MIN)
            nr_tries = NR_TRIES;
    }
    else
        nr_tries = NR_TRIES;
    

    vector <vector <double> > pqp_times;
    vector <int> pqp_attempts;

    // PQP Test
    test (nr_tries, pqp_times, pqp_attempts);

    std::cout << std::endl << std::endl;
    std::cout << "Analysis:" << std::endl;

    // Assume size of all 3 time vectors are the same
    for (size_t i = 0; i < pqp_times.size (); ++i)
    {
        double pqp_time = 0.0;
        for (size_t j = 0; j < pqp_times[i].size (); j++)
        {
            pqp_time += pqp_times[i][j];
        }

        std::cout << std::endl;
        if (i == 0)
            std::cout << " Apartment problem - Average Time (s)" << std::endl;
        else if (i == 1)
            std::cout << " Cubicles problem - Average Time (s)" << std::endl;
        else if (i == 2)
            std::cout << " \'Easy\' problem - Average Time (s)" << std::endl;

        std::cout << "             PQP: " << pqp_time / (double) nr_tries << "  " << nr_tries << "/" << pqp_attempts[i] << " successful" << std::endl;
    }


    return 0;
}

void test (int tries, vector <vector<double> > &times, vector<int> &attempts)
{
    std::cout << std::endl << "********************************************" << std::endl;
    std::cout << "Testing " << ("discrete") 
                 << " " << ("PQP") << " checker" << std::endl;

    vector <double> time;
    int successful = 0;
    int problem = 0;
    int nr_attempts = 0;

    while (problem != 3)
    {
        if (problem == 0) std::cout << std::endl << "Apartment problem" << std::endl;
        else if (problem == 1) std::cout << std::endl << "Cubicles problem" << std::endl;
        else if (problem == 2) std::cout << std::endl << "\'Easy\' problem" << std::endl;
        while (successful < tries)
        {
            // plan in SE3
            app::SE3RigidBodyPlanning setup;

            switch (problem)
            {
                case 0: configureApartmentProblem (setup); break;
                case 1: configureCubiclesProblem (setup); break;
                case 2: configureEasyProblem (setup); break;
                default: std::cout << "There is a problem" << std::endl; problem = 3; break;                
            }

            // setting collision checking resolution to 1% of the space extent
            setup.getSpaceInformation()->setStateValidityCheckingResolution(0.01);

            // use hPRM to plan
            setup.setPlanner (ob::PlannerPtr (new og::hPRM (setup.getSpaceInformation ())));
            setup.setup ();

            // Solve the problem.  Abort after 5.0 seconds
            setup.solve(5.0);

            // Retry if the planner failed, except for the continuous collision checker case.
            if (setup.haveExactSolutionPath ())
            {
                successful++;
                time.push_back (setup.getLastPlanComputationTime ());
            }
            nr_attempts++;
        }

        times.push_back (time);
        time.clear ();
        successful = 0; 
        problem++;
        attempts.push_back (nr_attempts);
    }
}

void configureApartmentProblem (app::SE3RigidBodyPlanning &setup)
{
    std::string problem = "Apartment";
    std::string environment_mesh = std::string(OMPLAPP_RESOURCE_DIR) + "/3D/" + problem + "_env.dae";
    std::string robot_mesh = std::string (OMPLAPP_RESOURCE_DIR) + "/3D/" + problem + "_robot.dae";

    // load the robot and the environment
    setup.setEnvironmentMesh(environment_mesh.c_str());
    setup.setRobotMesh(robot_mesh.c_str());

    // Bounds for Apartment environment
    ob::RealVectorBounds bounds(3);
    bounds.low[0] = -73.76;
    bounds.low[1] = -179.59;
    bounds.low[2] = -0.03;
    bounds.high[0] = 295.77;
    bounds.high[1] = 168.26;
    bounds.high[2] = 90.39;

    // Start/Goal pair for the Apartment environment
    ob::ScopedState<ob::SE3StateSpace> start (setup.getSpaceInformation ());
    start->setX (-31.19);
    start->setY (-99.85);
    start->setZ (36.46);
    start->rotation ().setIdentity ();

    ob::ScopedState<ob::SE3StateSpace> goal (setup.getSpaceInformation ());
    goal->setX (140.0);
    goal->setY (0.0);
    goal->setZ (36.46);
    goal->rotation ().setIdentity ();

    // Set start and goal
    setup.setStartAndGoalStates(start, goal);

    // Bound the state space
    setup.getSpaceInformation ()->getStateSpace ()->as <ob::SE3StateSpace> ()->setBounds (bounds);
}

void configureEasyProblem (app::SE3RigidBodyPlanning &setup)
{
    std::string problem = "Easy";
    std::string environment_mesh = std::string(OMPLAPP_RESOURCE_DIR) + "/3D/" + problem + "_env.dae";
    std::string robot_mesh = std::string (OMPLAPP_RESOURCE_DIR) + "/3D/" + problem + "_robot.dae";

    // load the robot and the environment
    setup.setEnvironmentMesh(environment_mesh.c_str());
    setup.setRobotMesh(robot_mesh.c_str());

    // Bounds for Easy/TwistyCool environment
    ob::RealVectorBounds bounds(3);
    bounds.low[0] = 14.46;
    bounds.low[1] = -24.25;
    bounds.low[2] = -504.86;
    bounds.high[0] = 457.96;
    bounds.high[1] = 321.25;
    bounds.high[2] = -72.86;

    // Start/Goal pair for the Easy/TwistyCool environment
    ob::ScopedState<ob::SE3StateSpace> start (setup.getSpaceInformation ());
    start->setX (270.4);
    start->setY (50.0);
    start->setZ (-406.82);
    start->rotation ().setIdentity ();

    ob::ScopedState<ob::SE3StateSpace> goal (setup.getSpaceInformation ());
    goal->setX (270.4);
    goal->setY (50.0);
    goal->setZ (-186.82);
    goal->rotation ().setIdentity ();

    // Set start and goal
    setup.setStartAndGoalStates(start, goal);

    // Bound the state space
    setup.getSpaceInformation ()->getStateSpace ()->as <ob::SE3StateSpace> ()->setBounds (bounds);
}

void configureCubiclesProblem (app::SE3RigidBodyPlanning &setup)
{
    std::string problem = "cubicles";
    std::string environment_mesh = std::string(OMPLAPP_RESOURCE_DIR) + "/3D/" + problem + "_env.dae";
    std::string robot_mesh = std::string (OMPLAPP_RESOURCE_DIR) + "/3D/" + problem + "_robot.dae";

    // load the robot and the environment
    setup.setEnvironmentMesh(environment_mesh.c_str());
    setup.setRobotMesh(robot_mesh.c_str());

    // Bounds for cubicles environment
    ob::RealVectorBounds bounds(3);
    bounds.low[0] = -508.88;
    bounds.low[1] = -230.13;
    bounds.low[2] = -123.75;
    bounds.high[0] = 319.62;
    bounds.high[1] = 531.87;
    bounds.high[2] = 101.00;

    // Start/Goal pair for cubicles environment
    ob::ScopedState<ob::SE3StateSpace> start (setup.getSpaceInformation ());
    start->setX (-4.96);
    start->setY (-40.62);
    start->setZ (70.57);
    start->rotation ().setIdentity ();

    ob::ScopedState<ob::SE3StateSpace> goal (setup.getSpaceInformation ());
    goal->setX (200.00);
    goal->setY (-40.62);
    goal->setZ (70.57);
    goal->rotation ().setIdentity ();

    // Set start and goal
    setup.setStartAndGoalStates(start, goal);

    // Bound the state space
    setup.getSpaceInformation ()->getStateSpace ()->as <ob::SE3StateSpace> ()->setBounds (bounds);
}

